package org.dllwh.milvus.v1;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.StrUtil;
import com.google.common.collect.Lists;
import io.milvus.client.MilvusServiceClient;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.grpc.SearchResultData;
import io.milvus.grpc.SearchResults;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.dml.SearchParam;
import io.milvus.response.SearchResultsWrapper;
import lombok.extern.slf4j.Slf4j;
import org.dllwh.milvus.config.MilvusConfig;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * 把今天最好的表现当作明天最新的起点．．～
 * <p>
 * Today the best performance as tomorrow the newest starter!
 *
 * @ClassName MilvusSearchHelper
 * @Description Milvus V1 数据搜索操作
 * @Author: <a href="mailto:duleilewuhen@sina.com">独泪了无痕</a>
 * @Date 2025-05-24 14:40
 * @版本 V 1.0.1
 * @since JDK 17
 */
@Slf4j
public class MilvusSearchHelper {
    /**
     * 向量搜索 - 向量库中具体向量搜索
     *
     * @param collectName      指定集合名称
     * @param vectorFieldName  指定搜索的Field
     * @param metricType       指定度量类型（计算方式）
     * @param floatList        指定搜索向量
     * @param topK             指定返回多少条结果
     * @param outFieldList     指定返回所有字段，也可以指定部分字段列表
     * @param consistencyLevel 指定一致性级别
     * @param expr
     * @param params
     * @param partitionNames   如果只需要搜索某个分区的数据，则需要指定分区
     * @param <T>
     * @return
     */
    public static <T> List<T> search(String collectName, String vectorFieldName, MetricType metricType, List<Float> floatList, int topK,
                                     List<String> outFieldList, ConsistencyLevelEnum consistencyLevel, String expr, String params,
                                     List<String> partitionNames) {
        if (StrUtil.isBlank(collectName) || CollectionUtil.isEmpty(floatList)) {
            return null;
        }
        MilvusServiceClient milvusClient = MilvusConnectHelper.getMilvusServiceClient();

        // 构建搜索条件
        SearchParam.Builder searchParam = SearchParam.newBuilder()
                .withCollectionName(collectName)
                .withFloatVectors(Collections.singletonList(floatList))
                .withVectorFieldName(vectorFieldName)
                .withTopK(topK)
                .withMetricType(metricType)
                .withOutFields(outFieldList)
                .withConsistencyLevel(consistencyLevel)
                .withExpr(expr)
                .withParams(params);

        if (CollectionUtil.isNotEmpty(partitionNames)) {
            searchParam.withPartitionNames(partitionNames);
        }

        // 执行查询
        R<SearchResults> response = milvusClient.search(searchParam.build());

        int status = response.getStatus();

        if (R.Status.Success.getCode() != response.getStatus()) {
            log.error("search 失败( status = {} )", status);
            Exception exception = response.getException();
            if (exception != null) {
                String message = response.getMessage();
                log.info("search 失败，失败原因是：{}", message);
                return null;
            }
        }
        // 解析结果数据
        SearchResults searchResults = response.getData();
        if (searchResults == null) {
            return null;
        }
        SearchResultsWrapper searchResultsWrapper = new SearchResultsWrapper(searchResults.getResults());
        SearchResultData results = searchResults.getResults();
        int scoresCount = results.getScoresCount();
        List milvusResList = new ArrayList<>();
        return milvusResList;
    }

}
